import pandas as pd
from pathlib import Path
from datetime import datetime
import numpy as np
import re
import pdb
import arrow
from datetime import datetime, timedelta

# 计算量表评分结果


def calc_scale_score(df: pd.DataFrame):
    """计算量表评分结果

    Keyword arguments:
    df -- DataFrame

    Return: result -- 计算量表评分结果
    """
    # 获取列数
    num_columns = df.shape[1]

    # 计算量表评分结果
    result = df.iloc[:, 0:num_columns].apply(
        lambda x: (x[x >= 0] - 1).sum() if x.min() >= 0 else None, axis=1)

    return result


# 计算Beck自杀量表评分结果
def calc_suicide_scale_score(df: pd.DataFrame, col_ideas: int, col_risks: int):
    """计算自杀量表评分结果

    Keyword arguments:
    df -- DataFrame
    col_ideas -- int -- 自杀意念结束列
    col_risks -- int -- 自杀风险结束列

    Return: result -- 计算自杀量表评分结果
    """

    # 自杀意念的强度（最近一周）（偶数序号，0,2,4,6,8）
    ideas_odds_num = [x for x in range(
        0, col_ideas) if x % 2 == 0]

    # 自杀意念的强度（最消沉、最忧郁的时候）（奇数序号，1,3,5,7,9）
    ideas_evens_num = [x for x in range(
        0, col_ideas) if x % 2 != 0]

    # 自杀危险性（最近一周）（偶数序号，10,12....36）
    risks_odds_num = [x for x in range(
        col_ideas, col_risks) if x % 2 == 0]

    # 自杀危险性（最消沉、最忧郁的时候）（奇数序号，11,13...,37）
    risks_evens_num = [x for x in range(
        col_ideas, col_risks) if x % 2 != 0]

    result = []
    for i in range(len(df)):

        if all(df.iloc[i, col_ideas:col_risks] < 0):
            result.append("")
        else:
            # 自杀意念强度（最近一周）评分
            ideas_lastweek = df.iloc[i, ideas_odds_num].apply(
                lambda x: x - 1).sum()

            # 自杀意念的强度（最消沉、最忧郁的时候）评分
            ideas_depressed = df.iloc[i, ideas_evens_num].apply(
                lambda x: x - 1).sum()

            # 自杀危险性（最近一周）和自杀危险性（最消沉、最忧郁的时候）
            if all(df.iloc[i, col_ideas:col_risks] < 0):
                result.append(
                    [
                        {
                            "descriptions": "自杀意念强度（最近一周）",
                            "score": int(ideas_lastweek)},
                        {
                            "descriptions": "自杀意念强度（最消沉、最忧郁的时候）",
                            "score": int(ideas_depressed)
                        }
                    ]
                )
            else:
                # 自杀危险性（最近一周）评分
                risks_lastweek = df.iloc[i, risks_odds_num].apply(
                    lambda x: x - 1 if x < 4 else 0
                ).sum()

                # 自杀危险性（最消沉、最忧郁的时候）评分
                risks_depressed = df.iloc[i, risks_evens_num].apply(
                    lambda x: x - 1 if x < 4 else 0
                ).sum()

                result.append(
                    [
                        {
                            "descriptions": "自杀意念强度（最近一周）",
                            "score": int(ideas_lastweek)
                        },
                        {
                            "descriptions": "自杀意念强度（最消沉、最忧郁的时候）",
                            "score": int(ideas_depressed)
                        },
                        {
                            "descriptions": "自杀危险性（最近一周）",
                            "score": int(risks_lastweek)
                        },
                        {
                            "descriptions": "自杀危险性（最消沉、最忧郁的时候）",
                            "score": int(risks_depressed)
                        }
                    ]
                )
    return result

# 计算儿童创伤量表（CTQ）评分结果


def calc_childhood_trauma_score(df: pd.DataFrame,
                                col_emotional_abuse: list,
                                col_physical_abuse: list,
                                col_sexual_abuse: list,
                                col_emotional_neglect: list,
                                col_physical_neglect: list,):
    """计算儿童创伤量表评分结果

    Keyword arguments:
    df -- DataFrame
    col_emotional_abuse -- list -- 情感虐待列表
    col_physical_abuse -- list -- 躯体虐待列表
    col_sexual_abuse -- list -- 性虐待列表
    col_emotional_neglect -- list -- 情感忽视列表
    col_physical_neglect -- list -- 身体忽视列表

    Return: result -- 计算儿童创伤量表（CTQ）评分结果
    """
    # 获取列数
    num_columns = df.shape[1]

    # 初始化结果列表
    result = []

    col_emotional_abuse = [x - 1 for x in col_emotional_abuse]
    col_physical_abuse = [x - 1 for x in col_physical_abuse]
    col_sexual_abuse = [x - 1 for x in col_sexual_abuse]
    col_emotional_neglect = [x - 1 for x in col_emotional_neglect]
    col_physical_neglect = [x - 1 for x in col_physical_neglect]

    # 计算评分结果
    for i in range(len(df)):
        # 校验所有列的值是否小于0
        if all(df.iloc[i, 0:num_columns] > 0):
            # 计算情感虐待的评估结果
            emotional_abuse = df.iloc[i, col_emotional_abuse].apply(
                lambda x: x - 1
            ).sum()

            # 计算躯体虐待的评估结果
            physical_abuse = df.iloc[i, col_physical_abuse].apply(
                lambda x: x - 1
            ).sum()

            # 计算性虐待评估结果
            sexual_abuse = df.iloc[i, col_sexual_abuse].apply(
                lambda x: x - 1
            ).sum()

            # 计算情感忽视的评估结果
            emotional_neglect = df.iloc[i, col_emotional_neglect].apply(
                lambda x: x - 1
            ).sum()

            # 计算身体忽视的评估结果
            physical_neglect = df.iloc[i, col_physical_neglect].apply(
                lambda x: x - 1
            ).sum()

            # 计算结果
            result.append(
                [
                    {
                        "descriptions": "情感虐待",
                        "score": int(emotional_abuse)
                    },
                    {
                        "descriptions": "躯体虐待",
                        "score": int(physical_abuse)
                    },
                    {
                        "descriptions": "性虐待",
                        "score": int(sexual_abuse)
                    },
                    {
                        "descriptions": "情感忽视",
                        "score": int(emotional_neglect)
                    },
                    {
                        "descriptions": "身体忽视",
                        "score": int(physical_neglect)
                    }
                ]
            )
        else:
            result.append(None)

    return result

# 计算基线日期
def calc_baseline_time(first_evaluation_time: datetime.date, first_visit_period: str):
    """计算基线日期

    Keyword arguments:
    first_evaluation_time -- datetime.date -- 第一次量表评估时间
    first_visit_period -- str -- 第一次评估访视期

    Return: baseline_date
    """
    interval_set = {
        "基线": (0, 0, 0, 0),
        "2周":  (0, 0, 2, 0),
        "第28天":  (0, 0, 4, 0),
        "3个月":  (0, 3, 0, 0),
        "6个月":  (0, 6, 0, 0),
        "第1年末":  (1, 0, 0, 0),
        "第1年半":  (1, 6, 0, 0),
        "第2年末": (2, 0, 0, 0),
        "第2年半": (2, 6, 0, 0),
        "第3年末": (3, 0, 0, 0),
        "第3年半": (3, 6, 0, 0),
        "第4年末": (4, 0, 0, 0),
        "第4年半": (4, 6, 0, 0),
        "第5年末": (5, 0, 0, 0),
        "第5年半": (5, 6, 0, 0),

    }
    first_evaluation_time = arrow.get(first_evaluation_time)

    # 计算基线日期
    if interval_set.get(first_visit_period):
        # 计算基线日期
        baseline_date = first_evaluation_time.shift(
            years=-interval_set[first_visit_period][0],
            months=-interval_set[first_visit_period][1],
            weeks=-interval_set[first_visit_period][2],
            days=-interval_set[first_visit_period][3]
        ).date()

    else:
        baseline_date = None
    return baseline_date

# 访视期映射表单
def set_visit_period(days_difference: int):
    """计算访视期

    Keyword arguments:
    days_difference -- 天数差

    Return: visit_period
    """
    days_difference_range = [0, 1, 45, 120,
                             240, 450, 600, 810, 990, 1260, 1620]
    visit_periods = ['基线', '第28天', '3个月', '6个月',
                     '1年半', '2年半', '第2年末', '第3年末', '第4年末', '其他随访时间']

    # 查找小于等于天数差的最大边界
    lower_boundary = max(
        [x for x in days_difference_range if x <= days_difference])  # 下限

    index = days_difference_range.index(lower_boundary)

    visit_period = visit_periods[index]

    return visit_period

# 根据当前访视的编码计算访视期


def calc_visit_period(first_visit_period: int = 1, days_difference: int = 0):
    """计算访视期

    Keyword arguments:
    first_visit_period -- 访视期
    days_difference -- 天数差

    Return: visit_period
    """

    # 访视期相差天数映射
    days = [
        (1, 0),
        (2, 28),
        (3, 90),
        (4, 180),
        (5, 360),
        (7, 14),
        (8, 60),
    ]

    if any(day[0] == first_visit_period for day in days):
        day = [day for day in days if day[0] == first_visit_period]
        days_difference = days_difference + day[0][1]
        visit_period = set_visit_period(days_difference)

    else:
        visit_period = None

    return visit_period


# 根据基线时间和量表评估时间计算访视期
def calc_current_visit_period(baseline_time: datetime.date, current_evaluation_time_str: str):
    """计算访视期

    Keyword arguments:
    baseline_time -- datetime.date -- 基线时间
    current_evaluation_time_str -- str -- 量表评估时间

    Return: visit_period
    """
    # 格式化量表评估时间为datetime.date类型
    current_evaluation_time = datetime.strptime(
        current_evaluation_time_str, "%Y/%m/%d %H:%M:%S")

    # 计算量表评估时间和基线时间的天数差
    days_difference = (arrow.get(current_evaluation_time) -
                       arrow.get(baseline_time)).days

    # 访视期相差天数映射
    days_difference_range = [
        (0, 7, "基线"),
        (14, 7, "2周"),
        (28, 7, "第28天"),
        (60, 10, "2个月"),
        (90, 20, "3个月"),
        (180, 45, "6个月"),
        (360, 45, "第1年末"),
        (540, 45, "第1年半"),
        (740, 45, "第2年末"),
        (900, 45, "第2年半"),
        (1080, 45, "第3年末"),
        (1260, 45, "第3年半"),
        (1440, 45, "第4年末"),
        (1620, 45, "第4年半"),
        (1800, 45, "第5年末"),
        (1980, 45, "第5年半"),
    ]

    # 计算访视期
    if days_difference < 0:
        visit_period = "基线"
    else:
        for i in range(len(days_difference_range)):
            if days_difference >= days_difference_range[i][0] and days_difference < (days_difference_range[i+1][0] + days_difference_range[i][1]):
                visit_period = days_difference_range[i][2]
                break
            else:
                visit_period = "其他随访时间"
    return visit_period

# 格式化日期时间


def formate_date(date_str: str):
    """格式化日期

    Keyword arguments:
    date_str -- 日期字符串

    Return: date_str
    """
    date_str = date_str.replace('-', '/')

    # 指定时间字符串的格式
    if len(date_str) > 10:
        time_format = "%Y/%m/%d %H:%M:%S"
    else:
        time_format = "%Y/%m/%d"

    # 使用strptime()方法将字符串转换为datetime对象
    date_time_obj = datetime.strptime(date_str, time_format)

    # 使用date()方法从datetime对象中提取日期部分
    date_obj = date_time_obj.date()

    return date_obj


# 格式化手机号
def formate_phone(phone_number: str):
    """"格式化手机号

    Keyword arguments:
    phone_number -- 手机号

    Return: phone_number
    """
    if phone_number == -3:
        phone_number = None
    else:
        # 去掉空格和换行符
        phone_number = str(phone_number).replace(' ', '').replace('\n', '')
        # 去掉非数字字符
        phone_number = re.sub(r'\D', '', phone_number)
        # 去掉开头的0
        phone_number = re.sub(r'^0', '', phone_number)
        # 提取11位手机号
        phone_number = str(phone_number)[0:11]
    return phone_number


# 格式化研究编号
def formate_code(code: str):
    """格式化研究编号

    Keyword arguments:
    code -- 研究编号

    Return: code
    """
    # 使用正则表达式去除 "xx周" 或 "xx月"
    code = re.sub(r'(\(|\）)?(\d+(周|月))(\)|\）)', '', code)
    code = re.sub(r'[一二三四五六七八九十壹贰叁肆伍陆柒捌玖拾两](周|月)', '', code)
    code = re.sub(r'\d+周|\d+月', '', code)
    code = re.sub(r'[(（]?[\u4e00-\u9fa5]+[)）]?|\¥', '', code)
    code = re.sub(r'[)）,， ]', '', code)
    code = re.sub(r'[(（-]', '/', code)
    code = re.sub(r' ', '', code)
    code = re.sub(r'([A-Z]{1,2}\d+)([A-Z]{1,2}\d+)', r'\1/\2', code)

    # 统一转化成大写
    code = str(code).upper()

    # 提取访视期编码和研究编号
    pattern = re.compile(r'^([0-9]{1,2})([A-Z][0-9/A-Z]+)')
    match = re.match(pattern, code)
    if match:
        visit_period_code = match.group(1)
        research_code = match.group(2)
    else:
        visit_period_code = None
        research_code = code

    return visit_period_code, research_code

# 数据处理和清洗
def process_dataset(input_file_doctor: Path, input_file_patient: Path, sheet_name: str = 'Sheet1'):
    """数据处理和清洗

    Keyword arguments:
    input_file_doctor -- 输入文件(医生端)
    input_file_patient -- 输入文件(患者端)
    sheet_name -- 表单名称

    Return: df_doctor, df_patient
    """
    # 读取BDMDD医生端Excel文件并将其转化成DataFrame
    df_doctor = pd.read_excel(
        input_file_doctor, sheet_name=sheet_name, index_col=None)

    # 读取BDMDD患者端Excel文件并将其转化成DataFrame
    df_patient = pd.read_excel(
        input_file_patient, sheet_name=sheet_name, index_col=None)

    # 过滤数据
    df_patient = df_patient[
        (df_patient.iloc[:, 7] == 1) &  # 剔除掉不同意知情同意书的数据
        (~df_patient.iloc[:, 8].isin(['啊啊啊', '测试', "胡少华"]))  # 剔除掉测试数据(姓名列)
    ]

    # 选择第11列的数据作为手机号数据
    phone_datas_doctor = df_doctor.iloc[:, 11].apply(
        lambda x: formate_phone(x))
    phone_datas_patient = df_patient.iloc[:, 11].apply(
        lambda x: formate_phone(x))

    # 选择第9列的数据作为组别数据
    research_cohort_doctor = df_doctor.iloc[:, 9].map({
        1: 'MDD',
        2: 'BD'
    })
    research_cohort_patient = df_patient.iloc[:, 12].map({
        1: 'MDD',
        2: 'BD'
    })

    # 选择第10列的数据作为随访时间
    visit_period_doctor = df_doctor.iloc[:, 10].map({
        1: '基线',
        2: '第28天',
        3: '3个月',
        4: '6个月',
        5: '第1年末',
        6: '一年以上',
        7: '2周',
        8: '其他随访时间',
    })
    visit_period_patient = df_patient.iloc[:, 10].map({
        1: '基线',
        2: '',
    })

    # 选择第8列的数据作为姓名
    name_doctor = df_doctor.iloc[:, 8]
    name_patient = df_patient.iloc[:, 8]

    # 选择第7/9列的数据提取研究编号
    research_code_doctor = df_doctor.iloc[:, 7].apply(
        lambda x: formate_code(str(x))[1])
    research_code_patient = df_patient.iloc[:, 9].apply(
        lambda x: formate_code(str(x))[1])

    # 选择第7/9列的数据提取访视s编号
    visit_code_doctor = df_doctor.iloc[:, 7].apply(
        lambda x: formate_code(str(x))[0])
    visit_code_patient = df_patient.iloc[:, 9].apply(
        lambda x: formate_code(str(x))[0])

    # 选择第1列的数据入组的时间
    group_entry_time_doctor = np.where(
        df_doctor.iloc[:, 10] != 1, np.nan, pd.to_datetime(df_doctor.iloc[:, 1], format='mixed').dt.date)
    group_entry_time_patient = np.where(
        df_patient.iloc[:, 10] != 1, np.nan, pd.to_datetime(df_patient.iloc[:, 1], format='mixed').dt.date)

    # 患者端gad-7、phq-9、cudos、becks、ctq列数据处理
    gad_7 = calc_scale_score(
        df_patient.iloc[:, 30:37])

    phq_9 = calc_scale_score(
        df_patient.iloc[:, 37:46])

    cudos = calc_scale_score(
        df_patient.iloc[:, 46:59])

    beck = calc_suicide_scale_score(df_patient.iloc[:, 59:97],
                                    col_ideas=10,
                                    col_risks=38)

    ctq = calc_childhood_trauma_score(
        df_patient.iloc[:, 97:125],
        col_emotional_abuse=[3, 8, 14, 18, 25],
        col_physical_abuse=[9, 11, 12, 15, 17],
        col_sexual_abuse=[20, 21, 23, 24, 27],
        col_emotional_neglect=[5, 7, 13, 19, 28],
        col_physical_neglect=[1, 2, 4, 6, 26],
    )
    calc_scale_score(
        df_patient.iloc[:, 97:125])

    # 医生端HAMD-24、HAMA-14、MADRS、YMRS
    hamd = calc_scale_score(
        df_doctor.iloc[:, 70:94])
    hama = calc_scale_score(
        df_doctor.iloc[:, 94:108])
    madrs = calc_scale_score(
        df_doctor.iloc[:, 108:118])
    ymrs = calc_scale_score(
        df_doctor.iloc[:, 118:129])

    # 增加姓名列
    df_doctor.insert(7, '姓名', name_doctor)
    df_patient.insert(7, '姓名', name_patient)

    # 增加手机号列
    df_doctor.insert(8, '手机号', phone_datas_doctor)
    df_patient.insert(8, '手机号', phone_datas_patient)

    # 增加研究编号列
    df_doctor.insert(9, '研究编号', research_code_doctor)
    df_patient.insert(9, '研究编号', research_code_patient)

    # 增加研究组别列
    df_doctor.insert(10, '组别', research_cohort_doctor)
    df_patient.insert(10, '组别', research_cohort_patient)

    # 增加访视编码列
    df_doctor.insert(11, '访视编码', visit_code_doctor)
    df_patient.insert(11, '访视编码', visit_code_patient)

    # 增加随访时间列
    df_doctor.insert(12, '访视期', visit_period_doctor)
    df_patient.insert(12, '访视期', visit_period_patient)

    # 增加基线日期列
    df_doctor.insert(13, '基线日期', group_entry_time_doctor)
    df_patient.insert(13, '基线日期', group_entry_time_patient)

    # 医生表增加量表评分(HAMD-24、HAMA-14、MADRS、YMRS)结果列
    df_doctor.insert(14, 'HAMD', hamd)
    df_doctor.insert(15, 'HAMA', hama)
    df_doctor.insert(16, 'MADRS', madrs)
    df_doctor.insert(17, 'YMRS', ymrs)

    # 患者表增加量表评分(GAD-7、PHQ-9、CUDOS、Beck、CTQ)结果列
    df_patient.insert(14, 'GAD-7', gad_7)
    df_patient.insert(15, 'PHQ-9', phq_9)
    df_patient.insert(16, 'CUDOS', cudos)
    df_patient.insert(17, 'Beck', beck)
    df_patient.insert(18, 'CTQ', ctq)

    return df_doctor, df_patient


# 根据系统已存在的患者, 补全信息
def supplement_patient_data(df: pd.DataFrame, patient_data_file: Path):
    """根据系统已存在的患者,补全信息

    Keyword arguments:
    df -- 数据

    Return: df -- 返回数据
    """
    # 读取已存在的患者数据
    df_patient_data = pd.read_excel(patient_data_file, sheet_name='Sheet1')

    # 查询手机号手机号数据小于11位的数据
    df_empty_phone = df.loc[
        df["手机号"].astype(str).str.len() < 11
    ]

    # 遍历手机号为空的数据
    for index, row in df_empty_phone.iterrows():
        name = row['姓名']
        phone = row['手机号']
        code = row['研究编号']
        group = row['组别']
        visit_period = row['访视期']
        entry_time = row['基线日期']

        # 在已存在的患者数据中查找姓名相同的数据
        existing_patient_info: pd.DataFrame = df_patient_data.loc[
            (df_patient_data["姓名"] == name)
        ]
        # 如果查询到姓名相同且研究编号相识的数据
        if len(existing_patient_info) == 1 and (str(existing_patient_info["研究编号"].iloc[0]) in str(code) or str(code) in str(existing_patient_info["研究编号"].iloc[0])):
            # 补全手机号数据
            df.loc[index, "手机号"] = str(existing_patient_info["手机号"].iloc[0])

            # 补全组别数据
            if group is np.nan:
                df.loc[index, "组别"] = existing_patient_info['组别'].iloc[0]

            # 补全访视期和基线日期数据
            if visit_period not in ["基线", "第28天", "3个月", "6个月", "第1年末", "1年半", "第2年末", "2年半", "第3年末"]:
                # 根据基线日期和当前填表的时间计算访视期
                baseline_date = arrow.get(existing_patient_info['基线日期'].iloc[0], 'YYYY-MM-DD').date()   # 基线日期

                # 筛选出序号相同的数据
                scale_assessment_time = str(df["提交答卷时间"].loc[index])  # 填表时间

                # 补全访视期数据
                df.loc[index, "访视期"] = calc_current_visit_period(baseline_time=baseline_date,
                                                                 current_evaluation_time_str=scale_assessment_time)
                # 补全基线日期数据
                df.loc[index, "基线日期"] = datetime.strptime(
                    existing_patient_info['基线日期'].iloc[0][0:10], '%Y-%m-%d').date()

    # 查询手机号不为空且基线时间为空的数据
    df_empty_baseline = df.loc[
        df["手机号"].astype(str).str.len().eq(11) &
        df_empty_phone["基线日期"].isnull()
    ]

    # 查找手机号不为空且基线时间为空的数据不重复的手机号
    unique_phones = df_empty_baseline["手机号"].astype(str).unique()

    if len(unique_phones) > 0:
        # 遍历手机号列表
        for phone in unique_phones:
            # 查找手机号不为空且基线时间为空的数据
            df_phones = df.loc[
                (df["手机号"].astype(str) == phone)
            ]
            # 判断df_phones是否为空
            if not df_phones.empty:

                # 获取df_phone第一行的index
                phone_index = df_phones.index[0]

                # 获取手机号不为空且基线时间为空的数据的姓名
                name = df_phones.loc[phone_index, "姓名"]

                # 查找姓名和手机号相同的数据
                df_empty_baseline_phone_name = df_patient_data.loc[
                    (df_patient_data["姓名"] == name) &
                    (df_patient_data["手机号"].astype(str) == phone)
                ]

                # 判断df_empty_baseline_phone_name是否为空
                if not df_empty_baseline_phone_name.empty:
                    # 获取姓名和手机号相同的数据的基线日期,转化成日期格式
                    baseline_date = arrow.get(
                        str(df_empty_baseline_phone_name["基线日期"].iloc[0]), 'YYYY-MM-DD').date()

                    # 获取组别
                    group = df_empty_baseline_phone_name["组别"].iloc[0]

                    # 获取研究编号
                    code = df_empty_baseline_phone_name["研究编号"].iloc[0]

                    # 补全基线日期数据
                    df.loc[df_phones.index[0], "基线日期"] = baseline_date

                    for index, row in df_phones.iterrows():
                        # 计算访视期
                        visit_period = calc_current_visit_period(
                            baseline_time=baseline_date,
                            current_evaluation_time_str=row["提交答卷时间"])
                        # 补全访视期数据
                        df.loc[index, "访视期"] = visit_period

                        # 补全组别数据
                        if group is np.nan:
                            df.loc[index, "组别"] = df_empty_baseline_phone_name["组别"].iloc[0]

                        # 补全研究编号数据
                        if code is np.nan:
                            df.loc[index, "研究编号"] = df_empty_baseline_phone_name["研究编号"].iloc[0]

    return df

# 补全研究编码
def supplement_research_code(df: pd.DataFrame):
    """补全研究编码

    Keyword arguments:
    df -- DataFrame
    Return: df
    """
    # 修改错误数据
    df.loc[df['姓名'] == '庄宥析', '姓名'] = '庄宥淅'
    df.loc[df['研究编号'] == 'ＦＡ045', '研究编号'] = 'FA045'

    # 获取姓名列中不重复的姓名
    unique_names = df[df['姓名'].str.len() > 0]['姓名'].drop_duplicates().tolist()

    # 遍历姓名列表
    for name in unique_names:
        # 查找姓名列中包含name的行
        rows = df[df['姓名'] == name]
        if len(rows) > 1:
            # 提取research_code列中的数据
            research_codes = rows["研究编号"].apply(lambda x: str(x).split('/'))

            # 找出research_codes中包含最多元素的数组
            max_research_code = max(
                [item for item in research_codes if item != [""]], key=len)

            # 找出research_codes中包含max_research_code的3个元素
            max_research_code_same = [i[0:3] for i in max_research_code]

            # 将research_codes中包含max_research_code的元素替换为max_research_code的元素
            df.loc[df['姓名'] == name, "研究编号"] = research_codes.apply(
                lambda x: '/'.join(max_research_code) if str(x[0]) == '' or any(b[0:3] in max_research_code_same for b in x) else '/'.join(x))

    return df


# 补全手机号数据
def supplement_phone_data(df: pd.DataFrame):
    """补全手机号数据

    Keyword arguments:
    df -- DataFrame
    Return: df
    """
    # 获取姓名列中不重复的姓名
    unique_names = df[df['姓名'].str.len() > 0]['姓名'].drop_duplicates().tolist()

    # 遍历姓名列表
    for name in unique_names:
        # 查找姓名列中包含name的行
        rows = df[df['姓名'] == name]

        # 获取姓名相同列中不重复的研究编号
        unique_codes = rows[rows['研究编号'].str.len(
        ) > 0]["研究编号"].drop_duplicates().tolist()

        # 遍历研究编号列表
        for code in unique_codes:
            # 查找姓名和研究编号相同的行
            phones = df.loc[(df['姓名'] == name) & (
                df['研究编号'] == code) & (df['手机号'].notna())]

            if len(phones) > 0:
                # 如果手机号不为空，则将手机号赋值给姓名和研究编号相同的行
                df.loc[(df['姓名'] == name) & (df['研究编号'] == code),
                       '手机号'] = phones.iloc[-1]['手机号']

    return df

# 补全基线时间数据


def supplement_baseline_data(df: pd.DataFrame):
    """补全基线时间数据

    Keyword arguments:
    df -- DataFrame
    Return: df
    """
    # 获取列中不重复的手机号
    unique_phones = df[df['手机号'].str.len(
    ) > 0]['手机号'].drop_duplicates().tolist()

    # 遍历不重复的手机号列
    for phone in unique_phones:
        # 查找手机号列中包含phone的行
        rows = df[df['手机号'] == phone]

        # 查询访视期为基线的行
        baseline_rows = rows[rows['访视期'] == "基线"]

        # 如果基线日期不为空, 将访视期为基线的时间赋值给访视期为基线日期的行
        if not baseline_rows.empty:
            # 基线日期
            baseline_date = arrow.get(datetime.strptime(baseline_rows.iloc[0, 1], "%Y/%m/%d %H:%M:%S")).date()

            # 计算第一次评估访视期
            first_evaluation_date = rows.iloc[0]["访视期"]

            # 如果第一次访视不是基线，修改访视期为基线
            if first_evaluation_date != '基线':
                df.loc[rows.index[0], "访视期"] = "基线"

            # 补全基线日期数据
            df.loc[rows.index[0], "基线日期"] = baseline_date

        # 访视编码不为空、访视期和基线日期为空
        elif rows["访视期"].notnull().all() and rows["基线日期"].isna().all() and rows["访视编码"].notna().any():
            # 访视期
            set_visit_period = {
                "1": "基线",
                "2": "第28天",
                "3": "3个月",
                "4": "6个月",
                "5": "第1年末",
                "6": "第1年半",
                "7": "第2年末",
                "8": "第2年半",
                "9": "第3年末",
                "10": "第3年半",
                "11": "第4年末",
            }

            # 查询访视编码不为空的行
            for index, row in rows.loc[rows["访视编码"].notna()].iterrows():
                # 访视期
                visit_period = set_visit_period[row["访视编码"]]

                # 补全访视期数据
                df.loc[index, "访视期"] = visit_period
                rows.loc[index, "访视期"] = visit_period

            # 获取rows中有访视期的行
            rows_visit_period_notna = rows.loc[rows["访视编码"].notna()]

            # 获取rows中有访视期的行第一个访视期
            visit_period = rows_visit_period_notna["访视期"].iloc[0]

            # 获取rows中有访视期的行第一个评估日期
            evaluation_date = formate_date(
                date_str=rows_visit_period_notna["提交答卷时间"].iloc[0])

            # 计算基线时间
            baseline_date = calc_baseline_time(
                first_evaluation_time=evaluation_date,
                first_visit_period=visit_period)

            # 补全基线日期数据
            df.loc[rows.index[0], "基线日期"] = baseline_date

        else:
            # 第一次量表评估访视期
            first_visit_period = rows.iloc[0]["访视期"]

            # 第一次量表评估时间
            first_evaluation_time = formate_date(date_str=rows.iloc[0, 1])

            baseline_date = calc_baseline_time(
                first_evaluation_time=first_evaluation_time,
                first_visit_period=first_visit_period)

            # 补全基线时间数据
            df.loc[rows.index[0], "基线日期"] = baseline_date

    return df

# 补全访视期数据


def supplement_visit_period_data(df: pd.DataFrame):
    """补全访视期数据

    Keyword arguments:
    df -- DataFrame
    Return: df
    """
    # 筛选出访视编码不为空的数据
    rows = df.loc[(df["访视编码"].notna())]

    set_visit_period = {
        "1": "基线",
        "2": "第28天",
        "3": "3个月",
        "4": "6个月",
        "5": "第1年末",
        "6": "第1年半",
        "7": "第2年末",
        "8": "第2年半",
        "9": "第3年末",
        "10": "第3年半",
        "11": "第4年末",
    }
    # 遍历访视期数据
    for index, row in rows.iterrows():
        # 访视编码
        visit_period_code = row["访视编码"]

        # 匹配计算访视期
        visit_period = set_visit_period[visit_period_code]

        if row['访视期'] is not np.nan:
            # 更新访视期数据
            df.loc[index, "访视期"] = visit_period

    # 获取列中不重复的手机号
    unique_phones = df[df['手机号'].str.len(
    ) > 0]['手机号'].drop_duplicates().tolist()

    # 遍历不重复的手机号列
    for phone in unique_phones:
        # 查找手机号列中包含phone的行
        rows = df[df['手机号'] == phone]

        # 基线日期数据
        baseline_date_datas = rows[rows['基线日期'].notna()]

        # 访视期数据
        visit_period_datas = rows["访视期"]

        # 筛选访视期不为空的数据
        if len(rows) > 1 and not baseline_date_datas.empty:

            baseline_date = arrow.get(baseline_date_datas["基线日期"].iloc[0]).date()

            # 遍历手机号列中包含phone的行
            for index, row in rows.iterrows():
                if row["访视期"] in ["其他随访时间", "一年以上"] or len(row["访视期"]) == 0:
                    # 量表评估日期
                    evaluation_time = row.iat[1]

                    # 计算访视期
                    visit_period = calc_current_visit_period(
                        baseline_time=baseline_date,
                        current_evaluation_time_str=evaluation_time)

                    if visit_period == "基线":
                        # 查询第一个访视的日期
                        first_evaluation_date = formate_date(
                            date_str=rows["提交答卷时间"].iloc[0])

                        # 找到第一个值的索引
                        first_index = df[df['手机号'] == phone].index[0]

                        # 第一个访视更新基线日期数据
                        df.loc[first_index, "基线日期"] = first_evaluation_date

                        # 其他访视期的基线日期更新为NaN
                        df.loc[(df['手机号'] == phone) & (df.index != first_index),
                               "基线日期"] = np.nan

                        # 第一个访视更新访视期数据
                        df.loc[first_index, "访视期"] = "基线"

                        # 重新计算其他访视期
                        df.loc[(df['手机号'] == phone) & (df.index != first_index), "访视期"] = df.loc[(df['手机号'] == phone) & (df.index != first_index)].apply(lambda x: calc_current_visit_period(
                            baseline_time=first_evaluation_date,
                            current_evaluation_time_str=x["提交答卷时间"]), axis=1)

                    # 补全访视期数据
                    df.loc[index, "访视期"] = visit_period

    return df

# 补全组别数据


def supplement_group_data(df: pd.DataFrame):
    """补全组别数据

    Keyword arguments:
    df -- DataFrame
    Return: df
    """
    # 获取列中不重复的手机号
    unique_phones = df[df['手机号'].str.len(
    ) > 0]['手机号'].drop_duplicates().tolist()

    # 遍历不重复的手机号列
    for phone in unique_phones:
        # 查找手机号列中包含phone的行
        rows = df[df['手机号'] == phone]

        # 获取最后一个访视的组别
        group = rows[rows['组别'].notna()]['组别'].tail(
            1).item() if not rows['组别'].dropna().empty else None

        # 筛选组别为空的数据
        # if group and not rows['组别'].empty:
        if group and not rows[rows['组别'].isna()].empty:
    
            # 补全组别数据
            df.loc[(df['手机号'] == phone) & (df['组别'].isna()), '组别'] = group

    return df

# 医生和患者量表数据互相补充数据


def supplement_empty_data(df_source: pd.DataFrame, df_reference: pd.DataFrame):
    """医生和患者量表数据互相补充数据

    Keyword arguments:
    df_source -- 源数据
    df_reference -- 参考数据
    Return: df_doctor, df_patient
    """
    # 筛选手机号为空的数据
    empty_phone_source = df_source.loc[
        df_source.loc[:, "手机号"].isna()
    ]

    # 遍历研究编号列表,补全手机号数据
    for index, rows in empty_phone_source.iterrows():
        # 研究编号
        code = rows.loc["研究编号"]
        # 姓名
        name = rows.loc["姓名"]

        # 查找参照数据中包含code和name且手机号不为空的行
        rows_reference = df_reference.loc[
            (df_reference.loc[:, "研究编号"] == code) & (
                df_reference.loc[:, "姓名"] == name) & (df_reference.loc[:, "手机号"].notna())
        ]

        # 如果存在匹配的行
        if not rows_reference.empty:
            # 补全手机号
            df_source.loc[index, "手机号"] = rows_reference["手机号"].iloc[0]

            # 如果手机号相同的行数小于2且基线日期不为空，补全基线日期数据
            if len(df_source[df_source['手机号'] == rows_reference["手机号"].iloc[0]]) < 2:
                # 补全基线时间数据
                df_source.loc[index, "基线日期"] = rows_reference["基线日期"].iloc[0]

    # 筛选手机号不为空的数据
    not_empty_phone_source = df_source.loc[
        df_source.loc[:, "手机号"].notna()
    ]

    # 查询不重复的手机号
    unique_phones = not_empty_phone_source[not_empty_phone_source['手机号'].str.len(
    ) > 0]['手机号'].drop_duplicates().tolist()

    for phone in unique_phones:
        # 查找手机号列中包含phone的行
        rows = df_source[df_source['手机号'] == phone]

        # 筛选基线日期为空的数据
        empty_baseline_date_rows = rows[rows['基线日期'].isna()]

        # 如果存在基线日期为空的数据,补全基线日期的时间数据
        if not empty_baseline_date_rows.empty:
            # 筛选参考数据姓名和手机号相同的行
            rows_reference = df_reference.loc[
                (df_reference.loc[:, "手机号"] == phone) & (
                    df_reference.loc[:, "姓名"] == name)
            ]

            # 筛选源数据中访视期不为空且基线日期全部为空的数据
            rows_visit_period_notna = rows[rows['访视期'].notna() & (
                rows.loc[:, "基线日期"].isna()).all()]

            # 如果存在参考数据姓名和手机号相同的行
            if not rows_reference.empty:
                # 基线日期数据
                baseline_date = rows_reference.at[rows_reference.index[0], "基线日期"]
                # 补全基线日期数据
                df_source.at[df_source.index[df_source["手机号"] == phone][0],
                             "基线日期"] = baseline_date

            # 如果源数据中访视期不为空
            elif not rows_visit_period_notna.empty:
                # 获取第一行的数据
                first_row = rows_visit_period_notna.iloc[0]
                # 访问第一行量表评估时间
                evaluation_time = formate_date(
                    date_str=first_row['提交答卷时间'])
                # 访问第一行的访视期
                first_visit_period = first_row['访视期']
                # 计算基线时间
                baseline_date = calc_baseline_time(
                    first_evaluation_time=evaluation_time,
                    first_visit_period=first_visit_period)
                # 补全基线日期数据
                df_source.at[df_source.index[df_source["手机号"]
                                             == phone][0], "基线日期"] = baseline_date

            # 基线时间不为空的行
            else:
                # 获取基线日期
                baseline_date = rows[rows['基线日期'].notna()]['基线日期'].iloc[0]

                # 根据基线时间计算访视期
                for index, row in rows.iterrows():
                    # 访问量表评估时间
                    evaluation_time = row['提交答卷时间']

                    # 计算访视期
                    visit_period = calc_current_visit_period(
                        baseline_time=baseline_date,
                        current_evaluation_time_str=evaluation_time)

                    if visit_period == "基线":
                        # 查询第一个访视的日期
                        first_evaluation_date = formate_date(
                            date_str=rows["提交答卷时间"].iloc[0])

                        # 找到第一个值的索引
                        first_index = df_source[df_source['手机号']
                                                == phone].index[0]

                        # 第一个访视更新基线日期数据
                        df_source.loc[first_index,
                                      "基线日期"] = first_evaluation_date

                        # 其他访视期的基线日期更新为NaN
                        df_source.loc[(df_source['手机号'] == phone) & (df_source.index != first_index),
                                      "基线日期"] = np.nan

                        # 第一个访视更新访视期数据
                        df_source.loc[first_index, "访视期"] = "基线"

                        # 重新计算其他访视期
                        df_source.loc[(df_source['手机号'] == phone) & (df_source.index != first_index), "访视期"] = df_source.loc[(df_source['手机号'] == phone) & (df_source.index != first_index)].apply(lambda x: calc_current_visit_period(
                            baseline_time=first_evaluation_date,
                            current_evaluation_time_str=x["提交答卷时间"]), axis=1)

                    # 补全访视期数据
                    df_source.loc[index, "访视期"] = visit_period

    return df_source


# 统一组别
def uniform_group_data(df: pd.DataFrame):
    """统一组别

    Keyword arguments:
    df -- DataFrame
    Return: df
    """
    # 获取列中不重复的手机号
    unique_phones = df[df['手机号'].str.len(
    ) > 0]['手机号'].drop_duplicates().tolist()

    for phone in unique_phones:
        # 查找手机号列中包含phone的行
        rows = df[df['手机号'] == phone]

        if len(rows) > 1:
            # 获取rows最后一个不为空的组别数据
            last_group = rows[rows['组别'].notna(
            )]['组别'].tail(1).item() if not rows['组别'].dropna().empty else None

            # 遍历手机号列中包含phone的行
            for index, row in rows.iterrows():
                # 补全组别数据
                df.loc[index, "组别"] = last_group

    return df


# 主函数
if __name__ == '__main__':
    # 使用Path对象来构建文件路径，使其更具可移植性
    input_file_doctor = Path('input') / 'BDMDD医生端.xlsx'
    output_file_doctor = Path('output') / 'BDMDD医生端导出.xlsx'

    input_file_patient = Path('input') / 'MDDBD患者端.xlsx'
    output_file_patient = Path('output') / 'MDDBD患者端导出.xlsx'

    # 已存在患者数据文件的路径
    patient_data_file = Path('input') / '已存在患者数据.xlsx'

    # 数据处理和清洗
    dt = process_dataset(input_file_doctor=input_file_doctor,
                         input_file_patient=input_file_patient, sheet_name='Sheet1')

    # 补全研究编码
    dt_doctor = supplement_research_code(dt[0])
    dt_patient = supplement_research_code(dt[1])

    # 补全手机号数据
    dt_doctor = supplement_phone_data(dt_doctor)
    dt_patient = supplement_phone_data(dt_patient)

    # 补全基线日期
    dt_doctor = supplement_baseline_data(dt_doctor)
    dt_patient = supplement_baseline_data(dt_patient)

    # 补全访视期数据
    dt_doctor = supplement_visit_period_data(dt_doctor)
    dt_patient = supplement_visit_period_data(dt_patient)

    # 补全组别数据
    dt_doctor = supplement_group_data(dt_doctor)
    dt_patient = supplement_group_data(dt_patient)

    # 医生和患者量表数据互相补充数据
    dt_doctor = supplement_empty_data(dt_doctor, dt_patient)
    dt_patient = supplement_empty_data(dt_patient, dt_doctor)

    # 根据系统已存在的患者,补全信息
    dt_doctor = supplement_patient_data(dt_doctor, patient_data_file)  # 医生数据
    dt_patient = supplement_patient_data(dt_patient, patient_data_file)  # 患者数据

    # 统一组别数据(取最后一个不为空的组别数据)
    dt_doctor = uniform_group_data(dt_doctor)
    dt_patient = uniform_group_data(dt_patient)

    #  保存修改后的医生数据到Excel文件
    dt_doctor.to_excel(output_file_doctor, index=False, sheet_name='Sheet1')

    #  保存修改后的患者数据到Excel文件
    dt_patient.to_excel(output_file_patient, index=False, sheet_name='Sheet1')

    # 打印提示信息
    print("数据处理和清洗完成！")
